import tensorflow as tf
import argparse
import numpy as np
from gpt2_model import *
import json
import sys

sys.path.append("..")
from util.optimizer import *
from util.metrics import *

physical_devices = tf.config.list_physical_devices("GPU")
for gpu_instance in physical_devices:
    tf.config.experimental.set_memory_growth(gpu_instance, True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--json_file", type=str, default="../util/124M.json")
    args = parser.parse_args()

    print("load params start")
    with open(args.json_file, "r") as f:
        params = json.load(f)
    print("load params over")

    print("load data start")
    data = np.load("../chinese_data/text.npy")
    label = np.load("../chinese_data/label.npy")
    print("data shape", data.shape)
    state = np.random.get_state()
    np.random.shuffle(data)

    np.random.set_state(state)
    np.random.shuffle(label)
    print(data.shape)
    print("load data over")
    physical_devices = tf.config.list_physical_devices('GPU')
    tf.config.set_visible_devices(physical_devices[0:2], 'GPU')
    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
    BATCH_SIZE_PER_REPLICA = 8
    GLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
    train_data_batches = data
    train_label_batches = label
    with strategy.scope():
        model = GPTModel(params)
        model.build(input_shape=(None, 1))
        model.summary()
        opt = get_optimizer(params)
        verbose = 1

        model.compile(
            optimizer=opt,
            loss=[tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), None],
            metrics=[[PerplexityMetric()], []]
            # metrics=[[PerplexityMetric()],[]],
        )
    model.fit(
        train_data_batches,
        train_label_batches,
        batch_size=8*strategy.num_replicas_in_sync,
        epochs=5000,
        validation_split=params["eval_ratio"],
        verbose=verbose,
        callbacks=[
            tf.keras.callbacks.ModelCheckpoint(
                "../checkpoint_v2/checkpoint-{epoch}.hdf5",
                verbose=0,
                save_best_only=True,
                mode="min",
            )
        ],
    )

    model.save_weights(params["model_path"])